package cn.bugstack.deepseek;

import cn.bugstack.deepseek.executor.model.deepseek.config.DeepSeekConfig;
import cn.bugstack.deepseek.executor.parameter.ChatChoice;
import cn.bugstack.deepseek.executor.parameter.CompletionRequest;
import cn.bugstack.deepseek.executor.parameter.CompletionResponse;
import cn.bugstack.deepseek.executor.parameter.Message;
import cn.bugstack.deepseek.session.Configuration;
import cn.bugstack.deepseek.session.OpenAiSession;
import cn.bugstack.deepseek.session.OpenAiSessionFactory;
import cn.bugstack.deepseek.session.defaults.DefaultOpenAiSessionFactory;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.logging.HttpLoggingInterceptor;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.apache.commons.lang3.StringUtils;
import org.junit.Before;
import org.junit.Test;

import javax.annotation.Nullable;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

@Slf4j
public class DeepSeekApiTest {

    private OpenAiSession openAiSession;

    @Before
    public void test_config(){

        DeepSeekConfig deepSeekConfig = new DeepSeekConfig();
        deepSeekConfig.setApiKey("sk-3d267048d00a40e58a314132c2e10d76");

        Configuration configuration = new Configuration();
        configuration.setLevel(HttpLoggingInterceptor.Level.BODY);
        configuration.setDeepSeekConfig(deepSeekConfig);

        // 初始化会话工厂
        OpenAiSessionFactory factory = new DefaultOpenAiSessionFactory(configuration);
        // 开启会话
        this.openAiSession = factory.openSession();
    }

    @Test
    public void test_completions() throws Exception{

        // 1. 创建参数
        CompletionRequest request = CompletionRequest.builder()
                .stream(true)
                .messages(Collections.singletonList(Message.builder().role(CompletionRequest.Role.USER).content("你是谁，请介绍一下你自己").build()))
                .model(CompletionRequest.Model.DEEPSEEK_V3.getCode())
                .build();

        // 2. 请求等待
        CountDownLatch countDownLatch = new CountDownLatch(1);

        EventSource eventSource = openAiSession.completions(request, new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                if ("[DONE]".equalsIgnoreCase(data)) {
                    log.info("OpenAI 应答完成");
                    return;
                }

                CompletionResponse chatCompletionResponse = JSON.parseObject(data, CompletionResponse.class);
                List<ChatChoice> choices = chatCompletionResponse.getChoices();
                for (ChatChoice chatChoice : choices) {
                    Message delta = chatChoice.getDelta();
                    if (CompletionRequest.Role.ASSISTANT.getCode().equals(delta.getRole())) continue;

                    // 应答完成
                    String finishReason = chatChoice.getFinishReason();
                    if (StringUtils.isNoneBlank(finishReason) && "stop".equalsIgnoreCase(finishReason)) {
                        return;
                    }

                    log.info("测试结果：{}", delta.getContent());
                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                log.info("对话完成");
                countDownLatch.countDown();
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                log.info("对话异常");
                countDownLatch.countDown();
            }
        });
        countDownLatch.await();
    }



    /**
     * TODO 测试DeepSeek-R1模型，与V3区别是，返回的数据里多了一个reasoning_content：思维链内容，与 content 同级，需要实现逻辑展示这一字段
     * @throws Exception
     */
    @Test
    public void test_completions2() throws Exception{

        // 1. 创建参数
        CompletionRequest request = CompletionRequest.builder()
                .stream(false)
                .messages(Collections.singletonList(Message.builder().role(CompletionRequest.Role.USER).content("我亲戚家的孩子今年高考，想了解一下上海大学，能不能帮我介绍一下这个学校").build()))
                .model(CompletionRequest.Model.DEEPSEEK_R1.getCode())
                .build();

        // 2. 请求等待
        CountDownLatch countDownLatch = new CountDownLatch(1);

        EventSource eventSource = openAiSession.completions(request, new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                if ("[DONE]".equalsIgnoreCase(data)) {
                    log.info("OpenAI 应答完成");
                    return;
                }

                CompletionResponse chatCompletionResponse = JSON.parseObject(data, CompletionResponse.class);
                List<ChatChoice> choices = chatCompletionResponse.getChoices();
                for (ChatChoice chatChoice : choices) {
                    Message delta = chatChoice.getDelta();
                    if (CompletionRequest.Role.ASSISTANT.getCode().equals(delta.getRole())) continue;

                    // 如果有推理内容，先输出推理内容
                    String reasoningContent = delta.getReasoning_content();
                    if (StringUtils.isNotBlank(reasoningContent)) {
                        log.info("推理内容: {}", reasoningContent);
                    }

                    // 输出最终的内容
                    String content = delta.getContent();
                    if (StringUtils.isNotBlank(content)) {
                        log.info("测试结果：{}", content);
                    }

                    // 应答完成
                    String finishReason = chatChoice.getFinishReason();
                    if (StringUtils.isNoneBlank(finishReason) && "stop".equalsIgnoreCase(finishReason)) {
                        return;
                    }


                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                log.info("对话完成");
                countDownLatch.countDown();
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                log.info("对话异常");
                countDownLatch.countDown();
            }
        });
        countDownLatch.await();
    }

    /**
     * 测试同步调用返回对话结果，不使用SSE监听得到生成式内容
     * @throws Exception
     */
    @Test
    public void test_completions3() throws Exception {

        // 1. 创建参数
        CompletionRequest request = CompletionRequest.builder()
                .stream(true)
                .messages(Collections.singletonList(Message.builder().role(CompletionRequest.Role.USER).content("你是谁，请介绍一下你自己").build()))
                .model(CompletionRequest.Model.DEEPSEEK_V3.getCode())
                .build();

        CompletionResponse response = openAiSession.completions(request);

        response.getChoices().forEach(e -> {
            log.info("测试结果：{}", e.getMessage());
        });

    }
}
